# -*- coding: utf-8 -*-
"""
Created on Thu Nov 25 20:25:29 2021

@author: bwm
"""
import pandas as pd
class after_trade:
    #初始化现金、总市值、股票持仓、股票价值
    def __init__(self, me,my_history,start_date,find,cash):
        self.me = me
        self.my_history = my_history
        self.start_date = start_date
        self.find = find
        self.cash = cash
        
    def after_trade(self,date):
        stock_in_stock_hold = self.me.stock_hold.index.tolist()
        num = self.me.stock_hold['num'].tolist()
        price = []
        if len(stock_in_stock_hold) != 0:
            for stock in stock_in_stock_hold:
                price.append(self.find.get_data(date,stock,'收盘价'))
        stock_data = {'num':num,'price':price}
        self.me.stock_hold = pd.DataFrame(stock_data,index = stock_in_stock_hold)
        self.me.stock_value = sum(self.me.stock_hold['num']*self.me.stock_hold['price'])
        self.me.total_value = self.me.cash + self.me.stock_value
                
        #记录每日持仓
        self.my_history.hold_history.loc[date] = {'cash':self.me.cash,'stock':stock_in_stock_hold,'total_value':self.me.total_value}#记录持仓信息,可能做成stock_hold
            
        #记录收益率
        if date == self.start_date:#如果是第一天，收益率记为0
            return_today = 0
        else:
            return_today = self.me.total_value/self.cash -1#return_today = total_value(今天)/total-value(昨天)-1
        self.my_history.return_history.loc[date] = {'ratio':round(return_today,4)}#记录收益率,保留四位小数